import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.preprocessing import MinMaxScaler
import argparse  # 用于解析命令行参数

# 读取数据集
data = pd.read_csv('data\生猪价格数据集utf-8.csv')

# 选择需要预测的特征以及其他相关特征
features = ['生猪价格', 'M2', '社会消费品零售总额', '城镇居民可支配收入', '能繁母猪存栏', '养殖成本', '牛肉价格',
                   '白条鸡价格', '百度搜索指数', '申万行业指数——生猪养殖']
dataset = data[features].values

# print(dataset.shape)
# 数据归一化
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(dataset)
y_test_max = scaler.data_max_[0]  # 第一个特征是生猪价格的最大值
y_test_min = scaler.data_min_[0]  # 第一个特征是生猪价格的最小值

# 定义训练集和测试集大小
test_size = 50
train_size = len(scaled_data) - test_size

# 划分训练集和测试集
train_data = scaled_data[:train_size, :]
test_data = scaled_data[train_size:, :]


def fit(X_train, y_train, X_test, y_test, epochs=50, batch_size=32):
    # 构建模型
    model = tf.keras.Sequential([
        tf.keras.layers.LSTM(units=50, return_sequences=True, input_shape=(X_train.shape[1], X_train.shape[2])),
        tf.keras.layers.LSTM(units=50),
        tf.keras.layers.Dense(units=1)
    ])

    # 编译模型
    model.compile(optimizer='adam', loss='mean_squared_error')
    
    # 训练模型
    model.fit(X_train, y_train, epochs=50, batch_size=32, verbose=0)
    # 预测未来几年的生猪价格
    predictions = model.predict(X_test)
    # print(predictions)
    predictions = (predictions * (y_test_max - y_test_min)) + y_test_min

    # 反归一化测试集标签值
    y_test_unscaled = (y_test * (y_test_max - y_test_min)) + y_test_min

    return predictions, y_test_unscaled

    # 打印预测结果和实际值比较
    for i in range(len(predictions)):
        print(f"预测值：{predictions[i][0]}, 实际值：{y_test_unscaled[i]}")

# parser = argparse.ArgumentParser()
# parser.add_argument('timestep', type=int, help='timestep预测步长')
# args = parser.parse_args()
# timestep = args.timestep

timestep=10

X_train, y_train = [], []
for i in range(timestep, len(train_data)):
    X_train.append(train_data[i-timestep:i, :])
    y_train.append(train_data[i, 0])  # 第一个特征是生猪价格
X_train, y_train = np.array(X_train), np.array(y_train)

X_test, y_test = [], []
for i in range(timestep, len(test_data)):
    X_test.append(test_data[i-timestep:i, :])
    y_test.append(test_data[i, 0])  # 第一个特征是生猪价格
X_test, y_test = np.array(X_test), np.array(y_test)

predictions, y_test_unscaled = fit(X_train, y_train, X_test, y_test )
# print(type(predictions))
print(f"--------------------timestep:{timestep}-------------------, \n 预测值{predictions}, \n 原始值{y_test_unscaled}")

print(len(predictions), len(y_test_unscaled))

    # res = np.append(res, predictions)
    # pre_predictions = predictions
